import tkinter as tk
from tkinter import filedialog, ttk, messagebox
from PIL import Image, ImageTk
import os, time, pickle, cv2
from functools import wraps
import numpy as np
import cupy as cp

def timer(unit="s"): # 带参数的装饰器，可指定时间单位
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            start = time.perf_counter()
            result = func(*args, **kwargs)
            duration = time.perf_counter() - start
            if unit == "ms":
                print(f"函数 {func.__name__} 运行时间: {duration*1000:.2f} 毫秒")
            else:
                print(f"函数 {func.__name__} 运行时间: {duration:.4f} 秒")
            return result
        return wrapper
    return decorator

@timer()
def variance_along_axis_cp(frame_list):
    #frame_list_cp = cp.asarray(np.stack(frame_list, axis=0))
    # 沿第 0 轴计算方差
    #variance_along_axis_0 = cp.var(cp.asarray(np.stack(frame_list, axis=0)), axis=(1, 2))
    return cp.argmax(cp.var(frame_list, axis=(1, 2))).item()
    #print("沿第 0 轴计算方差", cp.argmax(variance_along_axis_0).item(), cp.max(variance_along_axis_0).item())

@timer()
def variance_along_axis_np(frame_list):
    # 沿第 0 轴计算方差
    variance_along_axis_0 = [] #np.var(frame_list, axis=0)
    for i in range(len(frame_list)):
        variance_along_axis_0.append(frame_list[i].var())
    # 打印结果的形状
    print("沿第 0 轴的方差形状:", len(variance_along_axis_0), variance_along_axis_0[0].shape)

class ExtractFrames:
    def __init__(self, max_index = 19):
        self.max_index = max_index
        self.frames = [None for i in range(max_index + 1)] #cp.zeros((200, 1080, 1920), dtype="uint8")
        self.frames_index = 0
    def add(self, frame):
        if self.frames_index > self.max_index:
            self.frames_index = 0
        self.frames[self.frames_index] = frame
        self.frames_index += 1
    @timer()
    def variance_argmax(self):
        frames = cp.zeros((200, 1080, 1920), dtype="uint8")
        for i in range(self.max_index + 1):
            gray = cv2.cvtColor(self.frames[i], cv2.COLOR_BGR2GRAY)
            laplacian = cv2.Laplacian(gray, cv2.CV_64F)
            frames[i] = cp.asarray(laplacian)
        variance = cp.var(frames, axis=(1, 2))
        return cp.argmax(variance).item(), cp.max(variance).item()


class PicklePy:
    def save_pkl(self, file_name, data):
        # 将对象序列化并保存到文件
        with open(file_name, 'wb') as file:
            pickle.dump(data, file)
    def load_pkl(self, file_name):
        # 从文件中加载并反序列化对象
        with open(file_name, 'rb') as file:
            return pickle.load(file)
PKL = PicklePy()

class VideoPlayer:
    def __init__(self, root):
        self.root = root
        self.root.title("Tkinter 视频播放器")
        
        # 视频相关变量
        self.video_path = ""
        self.cap = None
        self.playing = False
        self.total_frames = 0
        self.fps = 0
        self.current_frame = 0
        self.delay = 0
        self.screenshot_count = 0  # 截图计数器
        self.EXF = ExtractFrames()
        
        # 视频尺寸变量
        self.video_width = 800  # 默认宽度
        self.video_height = 600  # 默认高度
        self.keep_aspect_ratio = True  # 保持宽高比
        
        # 倍速播放相关变量
        self.speed_var = tk.DoubleVar(value=1.0)  # 当前播放速度
        self.speed_options = [0.5, 0.75, 1.0, 1.25, 1.5, 2.0, 3.0, 5.0]  # 可选速度
        
        # 创建UI
        self.create_widgets()
        
        # 绑定窗口关闭事件
        self.root.protocol("WM_DELETE_WINDOW", self.on_close)
    
    def create_widgets(self):
        # 视频显示区域
        self.video_label = tk.Label(self.root)
        self.video_label.pack(padx=10, pady=10)
        
        # 控制按钮区域
        control_frame = tk.Frame(self.root)
        control_frame.pack(pady=5)
        
        # 打开文件按钮
        self.open_btn = tk.Button(
            control_frame, 
            text="打开视频", 
            command=self.open_file,
            width=10
        )
        self.open_btn.pack(side=tk.LEFT, padx=5)
        
        # 后退5秒按钮
        self.rewind_btn = tk.Button(
            control_frame,
            text="←5秒",
            command=self.rewind_5_seconds,
            state=tk.DISABLED,
            width=10
        )
        self.rewind_btn.pack(side=tk.LEFT, padx=5)
        
        # 播放/暂停按钮
        self.play_btn = tk.Button(
            control_frame, 
            text="播放", 
            command=self.toggle_play,
            state=tk.DISABLED,
            width=10
        )
        self.play_btn.pack(side=tk.LEFT, padx=5)
        
        # 前进5秒按钮
        self.forward_btn = tk.Button(
            control_frame,
            text="5秒→",
            command=self.forward_5_seconds,
            state=tk.DISABLED,
            width=10
        )
        self.forward_btn.pack(side=tk.LEFT, padx=5)
        
        # 截图按钮
        self.screenshot_btn = tk.Button(
            control_frame,
            text="截图",
            command=self.save_screenshot,
            state=tk.DISABLED,
            width=10
        )
        self.screenshot_btn.pack(side=tk.LEFT, padx=5)
        
        # 倍速播放下拉菜单
        self.speed_menu = ttk.Combobox(
            control_frame, 
            textvariable=self.speed_var, 
            values=self.speed_options,
            state="readonly",
            width=5
        )
        self.speed_menu.pack(side=tk.LEFT, padx=5)
        self.speed_menu.bind("<<ComboboxSelected>>", self.on_speed_change)
        
        # 倍速播放标签
        speed_label = tk.Label(control_frame, text="倍速")
        speed_label.pack(side=tk.LEFT, padx=5)
        
        # 尺寸调整区域
        size_frame = tk.Frame(self.root)
        size_frame.pack(pady=5)
        
        # 尺寸标签
        size_label = tk.Label(size_frame, text="视频尺寸:")
        size_label.pack(side=tk.LEFT, padx=5)
        
        # 预设尺寸下拉菜单
        self.size_var = tk.StringVar()
        self.size_var.set("800x600")  # 默认值
        sizes = ["800x600", "640x480", "1024x768", "1280x720", "自定义"]
        self.size_menu = ttk.Combobox(size_frame, textvariable=self.size_var, values=sizes, state="readonly")
        self.size_menu.pack(side=tk.LEFT, padx=5)
        self.size_menu.bind("<<ComboboxSelected>>", self.on_size_change)
        
        # 自定义尺寸按钮
        self.custom_size_btn = tk.Button(
            size_frame,
            text="设置",
            command=self.set_custom_size,
            width=5
        )
        self.custom_size_btn.pack(side=tk.LEFT, padx=5)
        
        # 保持宽高比复选框
        self.aspect_var = tk.BooleanVar(value=True)
        self.aspect_check = tk.Checkbutton(
            size_frame,
            text="保持宽高比",
            variable=self.aspect_var,
            command=self.toggle_aspect_ratio
        )
        self.aspect_check.pack(side=tk.LEFT, padx=5)
        
        # 进度条
        self.progress = ttk.Scale(
            self.root, 
            from_=0, 
            to=100, 
            orient=tk.HORIZONTAL,
            command=self.on_progress_change
        )
        self.progress.pack(fill=tk.X, padx=10, pady=5)
        self.progress.bind("<ButtonRelease-1>", self.on_progress_click)
        
        # 时间显示
        self.time_label = tk.Label(self.root, text="00:00:00 / 00:00:00")
        self.time_label.pack(pady=5)
        
        # 更新进度条的定时器
        self.update_interval = 30  # 毫秒
        self.after_id = None
    
    def open_file(self):
        file_path = filedialog.askopenfilename(
            title="选择视频文件",
            filetypes=[("视频文件", "*.mp4 *.avi *.mkv *.mov *.flv"), ("所有文件", "*.*")]
        )
        
        if file_path:
            self.video_path = file_path
            self.cap = cv2.VideoCapture(file_path)
            
            if self.cap.isOpened():
                # 获取视频信息
                self.total_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
                self.fps = self.cap.get(cv2.CAP_PROP_FPS)
                self.delay = int(1000 / self.fps)
                self.screenshot_count = 0  # 重置截图计数器
                
                # 启用播放按钮、前进、后退和截图按钮
                self.play_btn.config(state=tk.NORMAL)
                self.forward_btn.config(state=tk.NORMAL)
                self.rewind_btn.config(state=tk.NORMAL)
                self.screenshot_btn.config(state=tk.NORMAL)
                
                # 显示第一帧
                ret, frame = self.cap.read()
                if ret:
                    self.current_frame = 0
                    self.show_frame(frame)
                    self.update_progress()
                
                # 开始定时更新
                self.schedule_update(self.current_frame)
    
    def toggle_play(self):
        if self.playing:
            self.pause_video()
        else:
            self.play_video()
    
    def play_video(self):
        if self.cap and not self.playing:
            self.playing = True
            self.play_btn.config(text="暂停")
            self.update_video()
    
    def pause_video(self):
        if self.playing:
            self.playing = False
            self.play_btn.config(text="播放")
            if self.after_id:
                self.root.after_cancel(self.after_id)
    
    def update_video(self):
        if self.playing and self.cap:
            # 计算帧间隔时间，考虑倍速因素
            frame_interval = int(1000 / (self.fps * self.speed_var.get()))
            
            # 跳过一些帧以实现倍速效果
            frames_to_skip = int(self.speed_var.get() - 1)
            if frames_to_skip > 0:
                for _ in range(frames_to_skip):
                    self.cap.read()
            
            ret, frame = self.cap.read()
            if ret:
                self.current_frame += 1
                self.show_frame(frame)
                self.update_progress()
                self.schedule_update(frame_interval)
            else:
                # 视频结束，回到开头
                self.cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
                self.current_frame = 0
                self.playing = False
                self.play_btn.config(text="播放")
    
    def schedule_update(self, interval):
        if self.playing:
            self.after_id = self.root.after(interval, self.update_video)
    
    def on_speed_change(self, event):
        # 当速度改变时，如果视频正在播放，则重新开始播放
        if self.playing:
            self.pause_video()
            self.play_video()
    
    def show_frame(self, frame):
        self.EXF.add(frame)
        # 转换颜色空间从BGR到RGB
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        
        # 获取原始视频尺寸
        orig_height, orig_width = frame.shape[:2]
        
        # 调整大小
        if self.keep_aspect_ratio:
            # 计算保持宽高比的缩放比例
            ratio = min(self.video_width/orig_width, self.video_height/orig_height)
            new_width = int(orig_width * ratio)
            new_height = int(orig_height * ratio)
        else:
            new_width = self.video_width
            new_height = self.video_height
        
        frame = cv2.resize(frame, (new_width, new_height))
        
        # 转换为PIL图像格式
        img = Image.fromarray(frame)
        imgtk = ImageTk.PhotoImage(image=img)
        
        # 显示图像
        self.video_label.imgtk = imgtk
        self.video_label.config(image=imgtk)

    def is_frame_sharp_cupy(frame, threshold=100):
        # 转换为灰度图
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    
        # 将数据从NumPy数组转换为CuPy数组
        gray_cupy = cp.asarray(gray)
    
        # 计算拉普拉斯变换
        laplacian = cv2.Laplacian(gray_cupy, cv2.CV_64F)
    
        # 计算方差
        variance = cp.var(laplacian)
    
        # 将结果转换回NumPy数组（如果需要）
        variance_np = cp.asnumpy(variance)
    
        # 返回是否清晰（方差大于阈值）
        return variance_np > threshold, variance_np

    def update_progress(self):
        if self.total_frames > 0:
            progress = (self.current_frame / self.total_frames) * 100
            self.progress.set(progress)
            
            # 更新时间显示
            current_time = self.current_frame / self.fps
            total_time = self.total_frames / self.fps
            current_time_str = self.format_time(current_time)
            total_time_str = self.format_time(total_time)
            self.time_label.config(text=f"{current_time_str} / {total_time_str}")
    
    def format_time(self, seconds):
        hours = int(seconds // 3600)
        minutes = int((seconds % 3600) // 60)
        seconds = int(seconds % 60)
        return f"{hours:02d}:{minutes:02d}:{seconds:02d}"
    
    def on_progress_change(self, value):
        if self.cap and not self.playing:
            # 只有当视频暂停时才允许拖动进度条
            position = float(value) / 100 * self.total_frames
            self.cap.set(cv2.CAP_PROP_POS_FRAMES, int(position))
            ret, frame = self.cap.read()
            if ret:
                self.current_frame = int(position)
                self.show_frame(frame)
    
    def on_progress_click(self, event):
        # 处理进度条点击事件
        if self.cap and not self.playing:
            x = self.progress.winfo_width()
            click_pos = self.progress.winfo_pointerx() - self.progress.winfo_rootx()
            if click_pos < 0:
                click_pos = 0
            if click_pos > x:
                click_pos = x
            
            progress = (click_pos / x) * 100
            self.progress.set(progress)
            self.on_progress_change(progress)
    
    def rewind_5_seconds(self):
        if self.cap and self.cap.isOpened():
            # 计算5秒对应的帧数
            frames_to_rewind = int(5 * self.fps)
            new_frame = max(0, self.current_frame - frames_to_rewind)
            
            # 跳转到指定帧
            self.cap.set(cv2.CAP_PROP_POS_FRAMES, new_frame)
            ret, frame = self.cap.read()
            if ret:
                self.current_frame = new_frame
                self.show_frame(frame)
                self.update_progress()
    
    def forward_5_seconds(self):
        if self.cap and self.cap.isOpened():
            # 计算5秒对应的帧数
            frames_to_forward = int(5 * self.fps)
            new_frame = min(self.total_frames - 1, self.current_frame + frames_to_forward)
            
            # 跳转到指定帧
            self.cap.set(cv2.CAP_PROP_POS_FRAMES, new_frame)
            ret, frame = self.cap.read()
            if ret:
                self.current_frame = new_frame
                self.show_frame(frame)
                self.update_progress()
                
                # 如果视频在播放中，继续播放
                if self.playing:
                    self.update_video()
    
    def on_size_change(self, event):
        selected_size = self.size_var.get()
        if selected_size == "自定义":
            return
        
        # 解析预设尺寸
        width_str, height_str = selected_size.split('x')
        self.video_width = int(width_str)
        self.video_height = int(height_str)
        
        # 重新显示当前帧
        if self.cap and self.cap.isOpened():
            ret, frame = self.cap.read()
            if ret:
                self.show_frame(frame)
    
    def set_custom_size(self):
        # 创建自定义尺寸对话框
        dialog = tk.Toplevel(self.root)
        dialog.title("自定义视频尺寸")
        
        # 宽度标签和输入框
        tk.Label(dialog, text="宽度:").grid(row=0, column=0, padx=5, pady=5)
        width_entry = tk.Entry(dialog)
        width_entry.grid(row=0, column=1, padx=5, pady=5)
        width_entry.insert(0, str(self.video_width))
        
        # 高度标签和输入框
        tk.Label(dialog, text="高度:").grid(row=1, column=0, padx=5, pady=5)
        height_entry = tk.Entry(dialog)
        height_entry.grid(row=1, column=1, padx=5, pady=5)
        height_entry.insert(0, str(self.video_height))
        
        # 确定按钮
        def apply_size():
            try:
                new_width = int(width_entry.get())
                new_height = int(height_entry.get())
                if new_width > 0 and new_height > 0:
                    self.video_width = new_width
                    self.video_height = new_height
                    self.size_var.set(f"{new_width}x{new_height}")
                    
                    # 重新显示当前帧
                    if self.cap and self.cap.isOpened():
                        ret, frame = self.cap.read()
                        if ret:
                            self.show_frame(frame)
                    dialog.destroy()
                else:
                    messagebox.showerror("错误", "宽度和高度必须大于0")
            except ValueError:
                messagebox.showerror("错误", "请输入有效的数字")
        
        tk.Button(dialog, text="确定", command=apply_size).grid(row=2, column=0, columnspan=2, pady=10)
    
    def toggle_aspect_ratio(self):
        # 重新显示当前帧以应用宽高比设置
        if self.cap and self.cap.isOpened():
            ret, frame = self.cap.read()
            if ret:
                self.show_frame(frame)
    
    def save_screenshot(self):
        if self.cap and self.cap.isOpened():
            # 获取当前帧
            ret, frame = self.cap.read()
            if ret:
                argmax, max = self.EXF.variance_argmax()
                print(argmax, max)
                # 转换颜色空间从BGR到RGB
                frame = cv2.cvtColor(self.EXF.frames[argmax], cv2.COLOR_BGR2RGB)
                
                # 计算当前时间位置
                current_time = self.current_frame / self.fps
                time_str = self.format_time(current_time)
                
                # 确保视频目录存在
                video_dir = os.path.dirname(self.video_path)
                if not video_dir:
                    video_dir = "."  # 如果视频在当前目录，使用当前目录
                
                # 生成文件名
                self.screenshot_count += 1
                filename = f"screenshot_{self.screenshot_count}_{time_str.replace(':', '')}.png"
                file_path = os.path.join(video_dir, filename)
                
                # 保存图像
                img = Image.fromarray(frame)
                img.save(file_path)
                messagebox.showinfo("截图保存", f"截图已自动保存到视频目录:\n{file_path}")
            else:
                messagebox.showerror("错误", "无法获取当前帧")
        else:
            messagebox.showerror("错误", "没有打开的视频文件")
    
    def on_close(self):
        self.pause_video()
        if self.cap:
            self.cap.release()
        self.root.destroy()

if __name__ == "__main__":
    root = tk.Tk()
    app = VideoPlayer(root)
    root.mainloop()